import sys

import cv2
import numpy
import os


def read_images(path, sz=None):
    c = 0
    X, y = [], []
    for dir_name, dir_names, file_names in os.walk(path):
        for subdir_name in dir_names:
            subject_path = os.path.join(dir_name, subdir_name)
            for filename in os.listdir(subject_path):
                try:
                    if filename == ".directory":
                        continue
                    filepath = os.path.join(subject_path, filename)
                    im = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)

                    # resize to given size(if given)
                    if sz is not None:
                        im = cv2.resize(im, (200, 200))

                    X.append(numpy.asarray(im, dtype=numpy.uint8))
                    y.append(c)
                except IOError as e:
                    print("I/O error({0}): ({1})".format(e, str(e)))
                except:
                    print("Unexpected error:", sys.exc_info()[0])
                    raise
            c += 1

    return [X, y]


def face_rec():
    names = ["haoran.zhan", "张三", "李四"]

    [X, y] = read_images("./data")
    y = numpy.asarray(y, dtype=numpy.int32)

    model = cv2.face.EigenFaceRecognizer_create()
    model.train(numpy.asarray(X), numpy.asarray(y))

    camera = cv2.VideoCapture(0)

    face_cascade = cv2.CascadeClassifier('./cascades/haarcascade_frontalface_default.xml')
    while True:
        read, img = camera.read()
        faces = face_cascade.detectMultiScale(img, 1.3, 5)
        for (x, y, w, h) in faces:
            img = cv2.rectangle(img, (x, y), (x + w, y + h), (255, 0, 0), 2)
            gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
            roi = gray[x:x + w, y:y + h]

            roi = cv2.resize(roi, (200, 200), interpolation=cv2.INTER_LINEAR)
            params = model.predict(roi)
            print("Label: %s, Confidence: %.2f" % (params[0], params[1]))
            cv2.putText(img=img, text=names[params[0]], org=(x, y-20), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1, color=255, thickness=2)

        cv2.imshow("camera", img)
        if cv2.waitKey(1000 // 12) & 0xFF == ord("q"):
            break
    cv2.destroyAllWindows()


if __name__ == '__main__':
    face_rec()
